import base64
import hashlib
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad


class CryptoAESForAuth:
    SECRET_KEY = 'ab2d9c4491d647c4ac4947969d23366298bdbd376dba163b91f2cdb21255bbb3'
    def derive_key_and_iv(self, password, salt, key_length=32, iv_length=16):
        d = d_i = b''
        while len(d) < key_length + iv_length:
            d_i = hashlib.md5(d_i + password.encode('utf-8') + salt).digest()
            d += d_i
        return d[:key_length], d[key_length:key_length+iv_length]

    def decrypt_cryptojs_aes(self, ciphertext_b64):
        ciphertext = base64.b64decode(ciphertext_b64)
        assert ciphertext[:8] == b'Salted__'
        salt = ciphertext[8:16]
        key, iv = self.derive_key_and_iv(self.SECRET_KEY, salt)
        cipher = AES.new(key, AES.MODE_CBC, iv)
        decrypted = unpad(cipher.decrypt(ciphertext[16:]), AES.block_size)
        return decrypted.decode('utf-8')

    def encrypt_cryptojs_aes(self, text):
        salt = AES.get_random_bytes(8)
        key, iv = self.derive_key_and_iv(self.SECRET_KEY, salt)
        cipher = AES.new(key, AES.MODE_CBC, iv)
        ciphertext = cipher.encrypt(pad(text.encode('utf-8'), AES.block_size))
        encrypted = b'Salted__' + salt + ciphertext
        return base64.b64encode(encrypted).decode('utf-8')

if __name__ == "__main__":
    crypto_obj = CryptoAESForAuth()
    encrypted = crypto_obj.encrypt_cryptojs_aes("hello world")
    print("加密后（CryptoJS兼容）>>:", encrypted)

    decrypted = crypto_obj.decrypt_cryptojs_aes(encrypted)
    print("解密后 >>:", decrypted)
